package com.jd.genie.agent.agent;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jd.genie.agent.dto.Memory;
import com.jd.genie.agent.dto.Message;
import com.jd.genie.agent.dto.tool.ToolCall;
import com.jd.genie.agent.enums.AgentState;
import com.jd.genie.agent.enums.RoleType;
import com.jd.genie.agent.llm.LLM;
import com.jd.genie.agent.printer.Printer;
import com.jd.genie.agent.tool.BaseTool;
import com.jd.genie.agent.tool.ToolCollection;
import com.jd.genie.agent.util.SpringContextHolder;
import com.jd.genie.config.GenieConfig;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.context.ApplicationContext;

import java.util.*;
import java.util.function.Supplier;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * ReAct代理 - 基于ReAct模式的智能代理
 */
@Data
@Slf4j
public abstract class ReActAgent {

    // 核心属性
    private String systemPrompt;
    private String nextStepPrompt;
    public ToolCollection availableTools = new ToolCollection("fuck");
    private Memory memory = new Memory();
    protected LLM llm;
    protected final AgentContext context;

    // 执行控制
    private AgentState state = AgentState.IDLE;

    private final int maxSteps;

    private int currentStep = 0;
    private int duplicateThreshold = 2;

    // emitter
    Printer printer;

    // digital employee prompt
    private String digitalEmployeePrompt;

    public ReActAgent(AgentContext context)
    {
        ApplicationContext applicationContext = SpringContextHolder.getApplicationContext();
        GenieConfig genieConfig = applicationContext.getBean(GenieConfig.class);
        maxSteps = genieConfig.getPlannerMaxSteps();
        this.context = context;
    }

    /**
     * 运行代理主循环
     *
     * @param think 思考过程 返回true表示需要act
     */
    final public String baseRun(String query, Supplier<Boolean> think, Supplier<String> act)
    {
        setState(AgentState.IDLE);

        if (!query.isEmpty()) {
            updateMemory(RoleType.USER, query, null);
        }

        List<String> results = new ArrayList<>();
        try {
            while (currentStep < maxSteps && state != AgentState.FINISHED) {
                currentStep++;
                log.info("{} {} Executing step {}/{}", context.getRequestId(), getClass().getSimpleName(), currentStep, maxSteps);

                String stepResult;
                if (think.get()) {
                    stepResult = act.get();
                } else {
                    stepResult = "Thinking complete - no action needed";
                }
                results.add(stepResult);
            }

            if (currentStep >= maxSteps) {
                currentStep = 0;
                state = AgentState.IDLE;
                results.add("Terminated: Reached max steps (" + maxSteps + ")");
            }
        } catch (Exception e) {
            state = AgentState.ERROR;
            throw e;
        }

        return results.isEmpty() ? "No steps executed" : results.get(results.size() - 1);
    }

    /**
     * 更新代理记忆
     */
    public void updateMemory(RoleType role, String content, String base64Image, Object... args)
    {
        Message message = switch (role) {
            case USER -> Message.userMessage(content, base64Image);
            case SYSTEM -> Message.systemMessage(content, base64Image);
            case ASSISTANT -> Message.assistantMessage(content, base64Image);
            case TOOL -> Message.toolMessage(content, (String) args[0], base64Image);
            default -> throw new IllegalArgumentException("Unsupported role type: " + role);
        };
        memory.addMessage(message);
    }

    public String executeTool(ToolCall command)
    {
        if (command == null || command.getFunction() == null || command.getFunction().getName() == null) {
            return "Error: Invalid function call format";
        }

        String name = command.getFunction().getName();
        try {
            // 解析参数
            ObjectMapper mapper = new ObjectMapper();
            Object args = mapper.readValue(command.getFunction().getArguments(), Object.class);

            // 执行工具
            Object result = availableTools.execute(name, args);
            log.info("{} execute tool: {} {} result {}", context.getRequestId(), name, args, result);
            // 格式化结果
            if (Objects.nonNull(result)) {
                return (String) result;
            }
        } catch (Exception e) {
            log.error("{} execute tool {} failed ", context.getRequestId(), name, e);
        }
        return "Tool" + name + " Error.";
    }

    /**
     * 并发执行多个工具调用命令并返回执行结果
     *
     * @param commands 工具调用命令列表
     * @return 返回工具执行结果映射，key为工具ID，value为执行结果
     */
    public Map<String, String> executeTools(List<ToolCall> commands)
    {
        Map<String, String> result = new HashMap<>();
        for (ToolCall tooCall : commands) {
            String toolResult = executeTool(tooCall);
            result.put(tooCall.getId(), toolResult);
        }
        return result;
    }

    public void generateDigitalEmployee(String task)
    {
        // 1、参数检查
        if (StringUtils.isEmpty(task)) {
            return;
        }
        try {
            // 2. 构建系统消息（提取为独立方法）
            String formattedPrompt = formatSystemPrompt(task);
            Message userMessage = Message.userMessage(formattedPrompt, null);

            // 3. 调用LLM并处理结果
            String llmResponse = getLlm().ask(
                    context,
                    Collections.singletonList(userMessage),
                    Collections.emptyList(),
                    false,
                    0.01);

            log.info("requestId: {} task:{} generateDigitalEmployee: {}", context.getRequestId(), task, llmResponse);
            JSONObject jsonObject = parseDigitalEmployee(llmResponse);
            if (jsonObject != null) {
                log.info("requestId:{} generateDigitalEmployee: {}", context.getRequestId(), jsonObject);
                context.getToolCollection().updateDigitalEmployee(jsonObject);
                context.getToolCollection().setCurrentTask(task);
                // 更新 availableTools 添加数字员工
                availableTools = context.getToolCollection();
            } else {
                log.error("requestId: {} generateDigitalEmployee failed", context.getRequestId());
            }

        } catch (Exception e) {
            log.error("requestId: {} in generateDigitalEmployee failed,", context.getRequestId(), e);
        }
    }

    // 解析数据员工大模型响应
    private JSONObject parseDigitalEmployee(String response)
    {
        /**
         * 格式一：
         *      ```json
         *      {
         *          "file_tool": "市场洞察专员"
         *      }
         *      ```
         * 格式二：
         *      {
         *          "file_tool": "市场洞察专员"
         *      }
         */
        if (StringUtils.isBlank(response)) {
            return null;
        }
        String jsonString = response;
        String regex = "```\\s*json([\\d\\D]+?)```";
        Pattern pattern = Pattern.compile(regex);
        Matcher matcher = pattern.matcher(response);
        if (matcher.find()) {
            String temp = matcher.group(1).trim();
            if (!jsonString.isEmpty()) {
                jsonString = temp;
            }
        }
        try {
            return JSON.parseObject(jsonString);
        } catch (Exception e) {
            log.error("requestId: {} in parseDigitalEmployee error:", context.getRequestId(), e);
            return null;
        }
    }

    // 提取系统提示格式化逻辑
    private String formatSystemPrompt(String task)
    {
        String digitalEmployeePrompt = getDigitalEmployeePrompt();

        StringBuilder toolPrompt = new StringBuilder();
        for (BaseTool tool : context.getToolCollection().getToolMap().values()) {
            toolPrompt.append(String.format("工具名：%s 工具描述：%s\n", tool.getName(), tool.getDescription()));
        }

        // 替换占位符
        return digitalEmployeePrompt
                .replace("{{task}}", task)
                .replace("{{ToolsDesc}}", toolPrompt.toString())
                .replace("{{query}}", context.getQuery());
    }

}